package com.sunday.common.mq.rocket.brave;

import brave.Span;
import brave.Tracer;
import brave.Tracing;
import brave.internal.Nullable;
import brave.messaging.MessagingRequest;
import brave.propagation.CurrentTraceContext;
import brave.propagation.TraceContext;
import brave.propagation.TraceContext.Extractor;
import brave.propagation.TraceContext.Injector;
import brave.propagation.TraceContextOrSamplingFlags;
import brave.sampler.SamplerFunction;
import lombok.extern.slf4j.Slf4j;
import org.apache.rocketmq.client.hook.SendMessageContext;
import org.apache.rocketmq.client.hook.SendMessageHook;
import org.apache.rocketmq.client.producer.SendStatus;
import org.apache.rocketmq.client.trace.TraceConstants;
import org.apache.rocketmq.common.message.Message;

import static brave.Span.Kind.PRODUCER;

/**
 * 消息发送埋点，进行brave链路构建
 *
 * @author zsy
 * @see brave.spring.rabbit.TracingMessagePostProcessor
 * @see org.apache.rocketmq.client.trace.hook.SendMessageOpenTracingHookImpl
 * @since 2023/7/31
 */
@Slf4j
final class SendMessageBraveTracingHookImpl implements SendMessageHook {

    final RocketTracing rocketTracing;
    final Tracing tracing;
    final Tracer tracer;
    final CurrentTraceContext currentTraceContext;
    final Extractor<MessageProducerRequest> extractor;
    final SamplerFunction<MessagingRequest> sampler;
    final Injector<MessageProducerRequest> injector;
    @Nullable
    final String remoteServiceName;

    SendMessageBraveTracingHookImpl(RocketTracing springRabbitTracing) {
        this.rocketTracing = springRabbitTracing;
        this.tracing = springRabbitTracing.tracing;
        this.tracer = tracing.tracer();
        this.currentTraceContext = tracing.currentTraceContext();
        this.extractor = springRabbitTracing.producerExtractor;
        this.sampler = springRabbitTracing.producerSampler;
        this.injector = springRabbitTracing.producerInjector;
        this.remoteServiceName = springRabbitTracing.remoteServiceName;
    }

    @Override
    public String hookName() {
        return "SendMessageBraveTracingHook";
    }

    /**
     * 生产者消息前置处理
     *
     * @param context
     * @see brave.spring.rabbit.TracingMessagePostProcessor#postProcessMessage(org.springframework.amqp.core.Message)
     */
    @Override
    public void sendMessageBefore(SendMessageContext context) {
        if (context == null) {
            return;
        }
        Message msg = context.getMessage();
        MessageProducerRequest request = new MessageProducerRequest(msg);

        TraceContext maybeParent = currentTraceContext.get();
        // 与消息消费者不同，我们先尝试当前跨度，然后再尝试提取。这是正确的
        // 顺序，因为范围内的span应该优先于可能过时的头条目。
        // 注意:正确使用Brave检测不会导致报头条目过时，就像我们
        // 读取后总是清除消息头。
        Span span;
        if (maybeParent == null) {
            TraceContextOrSamplingFlags extracted =
                    rocketTracing.extractAndClearTraceIdHeaders(extractor, request, msg);
            span = rocketTracing.nextMessagingSpan(sampler, request, extracted);
        } else { // 如果我们在作用域中有一个span，假设header之前已经被清除了
            span = tracer.newChild(maybeParent);
        }
        // 如果是真的，则不会进行记录，也不会向zipkin报告任何内容。
        // 然而，这个跨度应该仍然被注入到传出请求中。使用此标志可以避免执行昂贵的操作 计算。
        if (!span.isNoop()) {
            setProducerSpan(span, msg, context);
            // 以隐式时间戳开始跨度。span可以在调用start之前修改。例如，您可以向span和添加标记在没有锁争用的情况下设置其名称
            span.start();
            context.setMqTraceContext(span);
        }
        //将链路信息注入到userProperties中
        injector.inject(span.context(), request);
    }

    void setProducerSpan(Span span, Message msg, SendMessageContext context) {
        span.kind(PRODUCER).name(TraceConstants.TO_PREFIX + msg.getTopic());
        if (remoteServiceName != null) span.remoteServiceName(remoteServiceName);
        maybeTag(span, RocketTracing.ROCKETMQ_TOPIC, msg.getTopic());
        maybeTag(span, TraceConstants.ROCKETMQ_TAGS, msg.getTags());
        maybeTag(span, TraceConstants.ROCKETMQ_KEYS, msg.getKeys());
        maybeTag(span, TraceConstants.ROCKETMQ_SOTRE_HOST, context.getBrokerAddr());
        maybeTag(span, TraceConstants.ROCKETMQ_MSG_TYPE, context.getMsgType().name());
        maybeTag(span, TraceConstants.ROCKETMQ_BODY_LENGTH, String.valueOf(msg.getBody().length));
    }

    void maybeTag(Span span, String tag, String value) {
        if (value != null) span.tag(tag, value);
    }

    @Override
    public void sendMessageAfter(SendMessageContext context) {
        if (context == null || context.getMqTraceContext() == null) {
            return;
        }
        if (context.getSendResult() == null) {
            return;
        }

        if (context.getSendResult().getRegionId() == null) {
            return;
        }

        Span span = (Span) context.getMqTraceContext();
        maybeTag(span, TraceConstants.ROCKETMQ_SUCCESS, String.valueOf(context.getSendResult().getSendStatus().equals(SendStatus.SEND_OK)));
        maybeTag(span, TraceConstants.ROCKETMQ_MSG_ID, context.getSendResult().getMsgId());
        maybeTag(span, TraceConstants.ROCKETMQ_REGION_ID, context.getSendResult().getRegionId());
        span.finish();
    }
}
